import sys
import time

sys.path.append('/home/zengxh/workspace/python_developer_tools') # 绝对路径
import argparse
import os
import shutil
import random
import glob
import copy
import numpy as np
import cv2
import multiprocessing

import pprint

import torch
import yaml
from imutils import paths
from sklearn.model_selection import train_test_split
from tqdm import tqdm

import lcnn
from dataset.constants import TB_DATATYPE, LR_DATATYPE
from demo import get_model_eval
from lcnn.config import C, M
from dataset.train_test_split import get_json_pts
from lcnn.models.line_vectorizer import LineVectorizer
from lcnn.models.multitask_learner import MultitaskHead, MultitaskLearner
from python_developer_tools.cv.utils.torch_utils import init_cudnn, init_seeds
from 预测图片并且生成json标准文件 import predict
from python_developer_tools.files.common import mkdir, resetDir


def reverse_pts(org_pts, type):
    reversed_pts = []
    if type == TB_DATATYPE:
        for pts in org_pts:
            if pts[0][0] > pts[1][0]:
                reversed_pts.append(pts[::-1])
            else:
                reversed_pts.append(pts)
    else:
        for pts in org_pts:
            if pts[0][1] > pts[1][1]:
                reversed_pts.append(pts[::-1])
            else:
                reversed_pts.append(pts)

    return reversed_pts

def calc_diff(first_pts, second_pts, type):
    # 计算真实坐标值和预测坐标值差，first_pts-真实值，second_pts-预测值
    diff_list = []
    if type == TB_DATATYPE:
        for i in range(len(first_pts)):
            for j in range(len(first_pts[i])):
                diff_list.append(abs(first_pts[i][j][1] - second_pts[i][j][0]))
    else:
        for i in range(len(first_pts)):
            for j in range(len(first_pts[i])):
                diff_list.append(abs(first_pts[i][j][0] - second_pts[i][j][1]))

    return diff_list

def get_avg_precision(diff_list, tol=0.1):
    # 计算准确率
    avg_precision = 0.0
    sum_precision = 0.0
    for diff in diff_list:
        accuarcy = 1.0 - tol * diff
        if accuarcy > 0:
            sum_precision += accuarcy

    return sum_precision / len(diff_list)

def predict_line_pts(img_path, edge_type,device,model):
    image = cv2.imread(img_path)
    imageHeight, imageWidth, _ = image.shape
    datadict = {}
    datadict["img"] = image
    datadict["datatype"] = edge_type
    datadict["device"] = device
    result_lines = predict(datadict, model)
    return result_lines

def get_mean_avg_precision(predict_dir,config_file,devices,checkpoint_path):
    model, device = get_model_eval(config_file, devices, checkpoint_path)

    sum_avg_precision = 0.0
    sum_predict_time = 0.0
    min_avg_precision = 0.6 # 1-0.1*(6.7/2) 将准确率小于0.5mm的都输出
    image_paths = list(paths.list_images(predict_dir))
    len_imgs = len(image_paths)

    for img_path in tqdm(image_paths):
        line_type = TB_DATATYPE
        json_path = img_path.replace('.jpg', '.json')
        _, img_name = os.path.split(img_path)
        img_type = img_name.split('.')[0][-1]
        if img_type == 'l' or img_type == 'r':
            line_type = LR_DATATYPE

        # 获取真实标注数据点
        true_pts = get_json_pts(json_path)
        true_pts = reverse_pts(true_pts, line_type)
        if len(true_pts) < 2:
            len_imgs = len_imgs - 1
            continue

        # 获取预测数据点
        start = time.time()
        pred_line_pts = predict_line_pts(img_path, line_type,device,model)
        sum_predict_time += time.time()-start

        if line_type == TB_DATATYPE:
            pred_line_pts = pred_line_pts[::-1]
        if len(pred_line_pts) < 2:
            continue

        # 计算真实点坐标值和预测点坐标值差异
        res_diff_list = calc_diff(true_pts, pred_line_pts, line_type)

        # 计算准确率
        single_avg_precision = get_avg_precision(res_diff_list)
        # 如果预测错误那么将预测错的图片提取出来
        if min_avg_precision > single_avg_precision:
            print('The accuracy of {} is: {:.2%}'.format(img_path, single_avg_precision))
            if TB_DATATYPE in img_path:
                noerror_dir = os.path.join(opt.predict_error_dir,TB_DATATYPE)
            else:
                noerror_dir = os.path.join(opt.predict_error_dir, LR_DATATYPE)
            mkdir(noerror_dir)
            shutil.copy(img_path, noerror_dir)
            shutil.copy(json_path, noerror_dir)
            originimg = cv2.imread(img_path)
            for (a, b) in pred_line_pts:
                _ = cv2.line(originimg, (int(a[1]), int(a[0])), (int(b[1]), int(b[0])), (0, 255, 0), thickness=2)
            cv2.imwrite(os.path.join(noerror_dir,"Y"+img_name),originimg)

        sum_avg_precision += single_avg_precision

    mean_avg_precision = sum_avg_precision / len_imgs
    print('The mean_avg_precision is: {:.2%} mean time is {}'.format(mean_avg_precision,sum_predict_time/len_imgs))
    return mean_avg_precision

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="获取杭州人工认为有缺陷大图")
    parser.add_argument('--devices',
                        default=r"0",
                        help="没有分的文件夹")
    parser.add_argument('--config_file',
                        default=r"config/wireframe.yaml",
                        help="没有分的文件夹")
    parser.add_argument('--checkpoint_dir',
                        default=r"/home/zengxh/workspace/lcnn/logs/210810-230633-88f281a-baseline/*/*/*.pth",
                        help="没有分的文件夹")
    parser.add_argument('--predict_dir',
                        default=r"/home/zengxh/medias/data/ext/creepageDistance/lab_datasets/",
                        help="没有分的文件夹")
    parser.add_argument('--predict_error_dir',
                        default=r"/home/zengxh/datasets/predict_error",
                        help="没有分的文件夹")
    opt = parser.parse_args()
    config_file = opt.config_file
    best_mean_avg_precision = 0
    best_model_path = ""
    resetDir(opt.predict_error_dir)

    # model_list=glob.glob(opt.checkpoint_dir)
    model_list = ["/home/zengxh/workspace/lcnn/logs/210902-043354-88f281a-baseline/checkpoint_best_swa.pth"]
    for model_path in model_list:
        mean_avg_precision = get_mean_avg_precision(opt.predict_dir, config_file, opt.devices, model_path)
        if best_mean_avg_precision < mean_avg_precision:
            best_mean_avg_precision= mean_avg_precision
            best_model_path = model_path
            print(best_model_path,best_mean_avg_precision)
    print(best_model_path, best_mean_avg_precision)


